import torch
from transformers import BertModel
from torchcrf import CRF


class BertCRF(torch.nn.Module):
    def __init__(self, config):
        super(BertCRF, self).__init__()
        self.bert = BertModel.from_pretrained('bert-base-chinese', num_labels=config.num_labels)
        self.dropout = torch.nn.Dropout(p=config.dropout)
        self.classifier = torch.nn.Linear(config.hidden_size, config.num_labels)
        self.crf = CRF(config.num_labels, batch_first=True)
        self.init_crf()

    def forward(self, input_ids):
        outputs = self.bert(input_ids=input_ids)
        logits = self.classifier(self.dropout(outputs[0]))
        return self.crf.decode(logits)

    def get_loss(self, input_ids, token_type_ids, attention_mask, labels):
        bert_outputs = self.bert(input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask)
        logits = self.classifier(self.dropout(bert_outputs[0]))
        loss = self.crf(logits, labels) * (-1)
        return self.crf.decode(logits), loss

    def init_crf(self):
        for p in self.crf.parameters():
            _ = torch.nn.init.uniform_(p, -1, 1)
